import numpy as np


class GomokuEnv:
    def __init__(self, size=15):
        self.size = size
        self.board = np.zeros((size, size), dtype=int)
        self.done = False
        self.winner = None

    def reset(self):
        self.board.fill(0)
        self.done = False
        self.winner = None
        return self.board

    def step(self, action, player):
        if self.done:
            raise ValueError("Game is over")
        x, y = action
        if self.board[x, y] != 0:
            raise ValueError("Invalid action")
        self.board[x, y] = player
        self.done, self.winner = self.check_winner()
        return self.board, self.done, self.winner

    def check_winner(self):
        # Implement check for winner logic
        # 检查是否有五子连线
        def count_consecutive(board, x, y, dx, dy, player):
            count = 0
            for i in range(-4, 5):
                nx, ny = x + i * dx, y + i * dy
                if (
                    0 <= nx < self.size
                    and 0 <= ny < self.size
                    and board[nx, ny] == player
                ):
                    count += 1
                    if count == 5:
                        return True
                else:
                    count = 0
            return False

        directions = [(1, 0), (0, 1), (1, 1), (1, -1)]
        for dx, dy in directions:
            if count_consecutive(self.board, x, y, dx, dy, player):
                return True, player

        if not np.any(self.board == 0):
            return True, 0  # 平局

        return False, None
